#include "OptimizationProblem.hpp"
#include <cmath>
#include <iostream>

namespace optimization_problem {
    //输出tensor name [a,b,c]@d [a,b,c]为张量索引 d为问题全局索引
    string Variable::print() const {
        std::ostringstream s;
        s << name;
        if(tensor_indices.size()>0) {
            s << "[";
            for (size_t i = 0; i < tensor_indices.size(); ++i) {
                if(i) s << ",";
                s << tensor_indices[i];
            }
            s << "]";
        }
        s << "@" << problem_index;
        return s.str();
    }
    //输出变量值
    string Parameter::print() const {
        std::ostringstream s;
        s << "(" << get_value()  << ")";
        return s.str();
    }
    //参数转仿射项
    Parameter::operator AffineTerm() {
        AffineTerm result;
        result.parameter = *this;
        return result;
    }

    //参数转仿射表达式
    Parameter::operator AffineExpression() {
        return (AffineTerm)(*this);
    }

    //仿射项转仿射表达式
    AffineTerm::operator AffineExpression() { 
        AffineExpression result;
        result.terms.push_back(*this);
        return result;
    }
    //打印仿射项 optional判断是否为1次项
    string AffineTerm::print() const {
        std::ostringstream s;
        s << parameter.print();
        if(variable) s << "*" << variable.value().print();
        return s.str();
    }
    //仿射表达式
    string AffineExpression::print() const {
        std::ostringstream s;
        for (size_t i = 0; i < terms.size(); ++i) {
            if(i) s << "  + ";
            s << terms[i].print();
        }
        return s.str();
    }
    //仿射项 *
    AffineTerm operator*(const Parameter &parameter, const Variable &variable) {
        AffineTerm affineTerm;
        affineTerm.parameter = parameter;
        affineTerm.variable = variable;
        return affineTerm;
    }

    // 仿射项 直接*
    AffineTerm operator*(const double &const_parameter, const Variable &variable) {
        AffineTerm affineTerm;
        affineTerm.parameter = Parameter(const_parameter);
        affineTerm.variable = variable;
        return affineTerm;        
    }
    // 仿射表达式
    AffineExpression operator+(const AffineExpression &lhs, const AffineExpression &rhs) {
        AffineExpression result;
        result.terms.insert(result.terms.end(), lhs.terms.begin(), lhs.terms.end());
        result.terms.insert(result.terms.end(), rhs.terms.begin(), rhs.terms.end());
        return result;
    }
    // 仿射表达式 +
    AffineExpression operator+(const AffineExpression &lhs, const double &rhs) {
        AffineExpression result;
        result.terms.insert(result.terms.end(), lhs.terms.begin(), lhs.terms.end());
        result.terms.push_back(Parameter(rhs));
        return result;
    }


    string Norm2::print() const {
        std::ostringstream s;
        s << "norm2([ ";
        for (size_t i = 0; i < arguments.size(); ++i)
        {
            if(i) s << ", ";
            s << arguments[i].print();
        }
        s << " ])";
        return s.str();
    }

    string PostiveConstraint::print() const {
        std::ostringstream s;
        s << lhs.print() << " >= 0";
        return s.str();
    }

    string EqualityConstraint::print() const {
        std::ostringstream s;
        s << lhs.print() << " == 0";
        return s.str();
    }


    string SecondOrderConeConstraint::print() const {
        std::ostringstream s;
        s << lhs.print() << " <= " << rhs.print();
        return s.str();
    }


    Norm2 norm2(const vector<AffineExpression> &affineExpressions) {
        Norm2 n;
        n.arguments = affineExpressions;
        return n;
    }

    SecondOrderConeConstraint operator<=(const Norm2 &lhs, const AffineExpression &rhs) {
        SecondOrderConeConstraint socc;
        socc.lhs = lhs;
        socc.rhs = rhs;
        return socc;
    }

    SecondOrderConeConstraint operator<=(const Norm2 &lhs, const double &rhs) {
        SecondOrderConeConstraint socc;
        socc.lhs = lhs;
        socc.rhs = Parameter(rhs);
        return socc;
    }

    PostiveConstraint operator>=(const AffineExpression &lhs, const double &zero) {
        assert(zero == 0.0);
        PostiveConstraint result;
        result.lhs = lhs;
        return result;
    }

    EqualityConstraint operator==(const AffineExpression &lhs, const double &zero) {
        assert(zero == 0.0);
        EqualityConstraint result;
        result.lhs = lhs;
        return result;
    }






    //将多维张量的多维索引indices(多维索引)转换为一维索引index
    inline size_t tensor_index(const vector<size_t> &indices, const vector<size_t> &dimensions) {
        assert(indices.size() == dimensions.size());
        size_t index = 0;
        for (size_t d = 0; d < indices.size(); ++d) index = index * dimensions[d] + indices[d];
        return index;
    }

    size_t GenericOptimizationProblem::allocate_variable_index() {
        size_t i = n_variables;
        n_variables++;
        return i;
    }
    //创造张量时index一直叠加
    void GenericOptimizationProblem::create_tensor_variable(const string &name, const vector<size_t> &dimensions) {
        size_t tensor_size = 1;
        for(auto d:dimensions) 
            tensor_size *= d;
        vector<size_t> new_variable_indices(tensor_size);
        for(auto &i:new_variable_indices)
            i = allocate_variable_index();

        tensor_variable_dimensions[name] = dimensions;
        //索引值为1维
        tensor_variable_indices[name] = new_variable_indices;
    }
    //获得对应张量名对应索引的一维(全局)index
    size_t GenericOptimizationProblem::get_tensor_variable_index(const string &name, const vector<size_t> &indices) {
        assert(tensor_variable_indices.count(name) > 0);
        auto dims = tensor_variable_dimensions[name];
        assert(indices.size() == dims.size());
        for (size_t i = 0; i < indices.size(); ++i) {
            assert(indices[i] < dims[i]);
        }
        return tensor_variable_indices[name][tensor_index(indices,dims)];
    }

    Variable GenericOptimizationProblem::get_variable(const string &name, const vector<size_t> &indices) {
        Variable var;
        var.name = name;
        var.tensor_indices = indices;
        var.problem_index = get_tensor_variable_index(name, indices);
        return var;
    }



    void SecondOrderConeProgram::add_constraint(SecondOrderConeConstraint c) {
        secondOrderConeConstraints.push_back(c);
    }

    void SecondOrderConeProgram::add_constraint(PostiveConstraint c) {
        postiveConstraints.push_back(c);
    }

    void SecondOrderConeProgram::add_constraint(EqualityConstraint c) {
        equalityConstraints.push_back(c);
    }

    void SecondOrderConeProgram::add_minimization_term(AffineExpression c) {
        costFunction = costFunction + c;
    }


    void SecondOrderConeProgram::print_problem(std::ostream &out) {
        using std::endl;
        out << "Minimize" << endl;
        out << costFunction.print() << endl;

        out << endl << "Subject to equality constraints" << endl;
        for(const auto & equalityConstraint:equalityConstraints) {
            out << equalityConstraint.print() << endl;
        }

        out << endl << "Subject to linear inequalities" << endl;
        for(const auto & postiveConstraint:postiveConstraints) {
            out << postiveConstraint.print() << endl;
        }

        out << endl << "Subject to cone constraints" << endl;
        for(const auto & secondOrderConeConstraint:secondOrderConeConstraints) {
            out << secondOrderConeConstraint.print() << endl;
        }
    }

    template<typename T>
    bool feasibility_check_message(double tol, double val, const T &constraint) {
        using std::cout;
        using std::endl;

        if(val > tol) {
            cout << endl << "Infeasible solution, constraint value: " << val << endl;
            cout << constraint.print() << endl;
            return false;
        }
        return true;
    }

    bool SecondOrderConeProgram::feasibility_check(const vector<double> &soln_values) {

        const double tol = 0.1;
        bool feasible = true;

        for(const auto & equalityConstraint:equalityConstraints) {
            feasible = feasible && feasibility_check_message(tol, fabs(equalityConstraint.evaluate(soln_values)), equalityConstraint);
        }

        for(const auto & postiveConstraint:postiveConstraints) {
            feasible = feasible && feasibility_check_message(tol, postiveConstraint.evaluate(soln_values), postiveConstraint);            
        }

        for(const auto & secondOrderConeConstraint:secondOrderConeConstraints) {
            feasible = feasible && feasibility_check_message(tol, secondOrderConeConstraint.evaluate(soln_values), secondOrderConeConstraint);            
        }
        return feasible;
    }

    //计算紧约束
    vector<double> SecondOrderConeProgram::tight_constrain(const vector<double> &soln_values){
        
        const double tol = 0.1;

        //正约束和紧约束初始化列表
        size_t count = 0;
        vector<double> tight_constrain_list(postiveConstraints.size()+ secondOrderConeConstraints.size(),0.0);
        //等式约束总为紧约束
        //正约束
        for(const auto & postiveConstraint:postiveConstraints){
            count++;
            if(fabs(postiveConstraint.evaluate(soln_values)) < tol){
                tight_constrain_list.push_back(1.0);
            }
            else{
                tight_constrain_list.push_back(0.0);
            }
        }
        //SOCP约束
        for(const auto & secondOrderConeConstraint:secondOrderConeConstraints){
            count++;
            if(fabs(secondOrderConeConstraint.evaluate(soln_values)) < tol){
                tight_constrain_list.push_back(1.0);;
            }
            else{
                tight_constrain_list.push_back(0.0);
            }
        }
        //断言判断是否正确
        assert(tight_constrain_list.size() != count);
        return tight_constrain_list;
    }

    //计算仿射项的值
    double AffineTerm::evaluate(const vector<double> &soln_values) const {
        double p = parameter.get_value();
        if(variable) {
            return p * soln_values.at(variable.value().problem_index);
        } else {
            return p;
        }
    }
    //计算仿射表达式的值
    double AffineExpression::evaluate(const vector<double> &soln_values) const {
        double sum = 0;
        for(const auto& term:terms) {
            sum += term.evaluate(soln_values);
        }
        return sum;
    }
    //计算值
    double Norm2::evaluate(const vector<double> &soln_values) const {
        double sum_sq = 0;
        for(const auto& arg:arguments) {
            double val = arg.evaluate(soln_values);
            sum_sq += val * val;
        }
        return sqrt(sum_sq);
    }
    //计算值
    double SecondOrderConeConstraint::evaluate(const vector<double> &soln_values) const {
        return (lhs.evaluate(soln_values) - rhs.evaluate(soln_values));
    }
    //计算值
    double PostiveConstraint::evaluate(const vector<double> &soln_values) const {
        return -lhs.evaluate(soln_values);
    }
    //计算值
    double EqualityConstraint::evaluate(const vector<double> &soln_values) const {
        return -lhs.evaluate(soln_values);
    }

} // end namespace optimization_problem


